"""
Binary classification demo for diabetes using the NeuralNetwork class.

This script trains a neural network to predict whether a patient has no diabetes (0)
or prediabetes/diabetes (1) using the BRFSS2015 health indicators dataset.

See diabetes-explanation.md for detailed feature and preprocessing info.

To run this demo, use:
    python -m demos.diabetes.demo-diabetes
from the project root directory.
"""

import numpy as np
import gzip
import csv
from sklearn.metrics import accuracy_score, confusion_matrix

# Import neural network components
from neural_network import Layer, ReLUActivation, NeuralNetwork, LossSoftmaxCCE
from neural_network.activation_sigmoid import SigmoidActivation

# Settings
# Network architecture - using recommended sizes from plan document
hidden_sizes = [64, 64, 32]
activation_hidden = ReLUActivation  # Changed from Sigmoid to ReLU
epochs = 100000
learning_rate = 0.001
batch_size = 64
test_split = 0.2
random_seed = 42

# Data file path - use package-relative path
import os
data_path = os.path.join(os.path.dirname(__file__),
                         'diabetes_binary_5050split_health_indicators_BRFSS2015.csv.gz')

# Feature names in order (excluding target)
feature_names = [
    "HighBP", "HighChol", "CholCheck", "BMI", "Smoker", "Stroke", "HeartDiseaseorAttack",
    "PhysActivity", "Fruits", "Veggies", "HvyAlcoholConsump", "AnyHealthcare", "NoDocbcCost",
    "GenHlth", "MentHlth", "PhysHlth", "DiffWalk", "Sex", "Age", "Education", "Income"
]
target_name = "Diabetes_binary"

# 1. Load and parse the gzipped CSV
print("Loading and parsing data...")
with gzip.open(data_path, 'rt') as f:
    reader = csv.DictReader(f)
    data = [row for row in reader]

# 2. Convert to numpy arrays
X = np.array([[float(row[feat]) for feat in feature_names] for row in data])
y_raw = np.array([int(float(row[target_name])) for row in data])

# 3. Train/test split
np.random.seed(random_seed)
indices = np.arange(X.shape[0])
np.random.shuffle(indices)
split = int(X.shape[0] * (1 - test_split))
train_idx, test_idx = indices[:split], indices[split:]
X_train, X_test = X[train_idx], X[test_idx]
y_train_raw, y_test_raw = y_raw[train_idx], y_raw[test_idx]

# 4. Standardize continuous/ordinal features (fit on train, apply to both)
# Indices of features to standardize (see explanation.md)
standardize_idx = [3, 13, 14, 15, 18, 19, 20]  # BMI, GenHlth, MentHlth, PhysHlth, Age, Education, Income
means = X_train[:, standardize_idx].mean(axis=0)
stds = X_train[:, standardize_idx].std(axis=0)
stds[stds == 0] = 1.0  # Prevent division by zero

def standardize(X, idx, means, stds):
    Xc = X.copy()
    Xc[:, idx] = (Xc[:, idx] - means) / stds
    return Xc

X_train = standardize(X_train, standardize_idx, means, stds)
X_test = standardize(X_test, standardize_idx, means, stds)

# 5. One-hot encode target
y_train = np.eye(2)[y_train_raw].T  # shape (2, N)
y_test = np.eye(2)[y_test_raw].T

# 6. Transpose features for NN (shape: features x samples)
X_train = X_train.T
X_test = X_test.T

# 7. Build the network
input_size = X_train.shape[0]
# Initialize network with improved weight initialization
network = NeuralNetwork()
prev_size = input_size
for size in hidden_sizes:
    # Use Xavier/Glorot initialization for better performance with ReLU
    weights = np.random.randn(size, prev_size) * np.sqrt(2.0 / prev_size)
    biases = np.zeros((size, 1))  # Initialize biases to zero
    network.add_layer(Layer(prev_size, size, weights=weights, biases=biases))
    network.add_layer(activation_hidden())
    prev_size = size
# Output layer with proper initialization
weights = np.random.randn(2, prev_size) * np.sqrt(2.0 / prev_size)
biases = np.zeros((2, 1))
network.add_layer(Layer(prev_size, 2, weights=weights, biases=biases))  # 2 output classes

# 8. Loss function
loss_obj = LossSoftmaxCCE()

# Check class distribution in train and test sets
print("\nClass distribution:")
print(f"Train set - Class 0: {np.sum(y_train_raw == 0)}, Class 1: {np.sum(y_train_raw == 1)}")
print(f"Test set  - Class 0: {np.sum(y_test_raw == 0)}, Class 1: {np.sum(y_test_raw == 1)}")

# Examine initial predictions before training
print("\nInitial predictions (before training):")
initial_pred = network.forward_pass(X_test)
initial_pred_labels = np.argmax(initial_pred, axis=0)
initial_true_labels = np.argmax(y_test, axis=0)
initial_acc = np.mean(initial_pred_labels == initial_true_labels)

# Calculate initial confusion matrix
initial_cm = np.zeros((2, 2), dtype=int)
for true, pred in zip(initial_true_labels, initial_pred_labels):
    initial_cm[true, pred] += 1

print(f"Initial test accuracy: {initial_acc:.4f}")
print("Initial confusion matrix (rows: true, cols: pred):")
print(initial_cm)
print(f"Class 0 prediction rate: {np.sum(initial_pred_labels == 0) / len(initial_pred_labels):.4f}")
print(f"Class 1 prediction rate: {np.sum(initial_pred_labels == 1) / len(initial_pred_labels):.4f}")

# Print raw network outputs before training
print("\nRaw network outputs (first 5 samples):")
print(initial_pred[:, :5])

# Apply softmax to see probabilities
def softmax(x):
    exp_x = np.exp(x - np.max(x, axis=0, keepdims=True))
    return exp_x / np.sum(exp_x, axis=0, keepdims=True)

initial_probs = softmax(initial_pred[:, :5])
print("\nSoftmax probabilities (first 5 samples):")
print(initial_probs)
print("Class 0 probabilities:", initial_probs[0])
print("Class 1 probabilities:", initial_probs[1])

# 9. Training loop
print("\nTraining network...")

# Use the enhanced train method with accuracy metric
network.train(
    X_train, y_train,
    loss_obj,
    epochs,
    learning_rate,
    batch_size,
    shuffle=True,
    verbose=True,
    X_val=X_test,
    y_val=y_test,
    metric="accuracy",
    plot=True
)

# Calculate weight statistics for all layers
max_weight = -float('inf')
min_weight = float('inf')
max_abs_weight = 0

for layer in network.layers:
    if hasattr(layer, 'weights'):
        layer_max = np.max(layer.weights)
        layer_min = np.min(layer.weights)
        layer_max_abs = np.max(np.abs(layer.weights))
        
        max_weight = max(max_weight, layer_max)
        min_weight = min(min_weight, layer_min)
        max_abs_weight = max(max_abs_weight, layer_max_abs)

print(f"Final weight statistics - max_w: {max_weight:.4f} - min_w: {min_weight:.4f} - max_abs_w: {max_abs_weight:.4f}")

# 10. Final evaluation
print("\nFinal evaluation on test set:")
y_pred_test = network.forward_pass(X_test)
y_pred_labels = np.argmax(y_pred_test, axis=0)
y_true_labels = np.argmax(y_test, axis=0)
acc = accuracy_score(y_true_labels, y_pred_labels)
cm = confusion_matrix(y_true_labels, y_pred_labels)
print(f"Test accuracy: {acc:.4f}")

# Print a more detailed confusion matrix with labels
print("\nDetailed Confusion Matrix:")
print("                      Predicted")
print("                  No Diabetes  Diabetes")
print(f"Actual No Diabetes    {cm[0,0]:4d}      {cm[0,1]:4d}")
print(f"      Diabetes        {cm[1,0]:4d}      {cm[1,1]:4d}")
print("\nConfusion Matrix Interpretation:")
print(f"True Negatives (TN): {cm[0,0]} - Correctly predicted as No Diabetes")
print(f"False Positives (FP): {cm[0,1]} - Incorrectly predicted as Diabetes")
print(f"False Negatives (FN): {cm[1,0]} - Incorrectly predicted as No Diabetes")
print(f"True Positives (TP): {cm[1,1]} - Correctly predicted as Diabetes")



# Function to interpret feature values in human-readable form
# This is one big function just for making the data in a nice string
#  format. Nothing mathematical or technical is happening here.
def interpret_features(features, feature_names):
    interpretations = []
    
    # Create a dictionary for easier access
    feature_dict = {name: value for name, value in zip(feature_names, features)}
    
    # Binary features (0/1)
    binary_features = {
        "HighBP": ("High Blood Pressure", {0: "No", 1: "Yes"}),
        "HighChol": ("High Cholesterol", {0: "No", 1: "Yes"}),
        "CholCheck": ("Cholesterol Check in 5 Years", {0: "No", 1: "Yes"}),
        "Smoker": ("Smoker (100+ cigarettes)", {0: "No", 1: "Yes"}),
        "Stroke": ("Ever Had Stroke", {0: "No", 1: "Yes"}),
        "HeartDiseaseorAttack": ("Heart Disease or Attack", {0: "No", 1: "Yes"}),
        "PhysActivity": ("Physical Activity in Past 30 Days", {0: "No", 1: "Yes"}),
        "Fruits": ("Consume Fruit 1+ Times/Day", {0: "No", 1: "Yes"}),
        "Veggies": ("Consume Vegetables 1+ Times/Day", {0: "No", 1: "Yes"}),
        "HvyAlcoholConsump": ("Heavy Alcohol Consumption", {0: "No", 1: "Yes"}),
        "AnyHealthcare": ("Has Health Care Coverage", {0: "No", 1: "Yes"}),
        "NoDocbcCost": ("Couldn't See Doctor Due to Cost", {0: "No", 1: "Yes"}),
        "DiffWalk": ("Difficulty Walking/Climbing Stairs", {0: "No", 1: "Yes"}),
        "Sex": ("Sex", {0: "Female", 1: "Male"})
    }
    
    # Ordinal features
    ordinal_features = {
        "GenHlth": ("General Health", {
            1: "Excellent", 2: "Very Good", 3: "Good", 4: "Fair", 5: "Poor"
        }),
        "Age": ("Age Category", {
            1: "18-24", 2: "25-29", 3: "30-34", 4: "35-39", 5: "40-44",
            6: "45-49", 7: "50-54", 8: "55-59", 9: "60-64", 10: "65-69",
            11: "70-74", 12: "75-79", 13: "80+"
        }),
        "Education": ("Education Level", {
            1: "Never attended school", 2: "Elementary", 3: "Some high school",
            4: "High school graduate", 5: "Some college", 6: "College graduate"
        }),
        "Income": ("Income Level", {
            1: "< $10k", 2: "$10k-$15k", 3: "$15k-$20k", 4: "$20k-$25k",
            5: "$25k-$35k", 6: "$35k-$50k", 7: "$50k-$75k", 8: "$75k+"
        })
    }
    
    # Continuous features
    continuous_features = {
        "BMI": "Body Mass Index",
        "MentHlth": "Days Mental Health Not Good (Past 30 Days)",
        "PhysHlth": "Days Physical Health Not Good (Past 30 Days)"
    }
    
    # Process binary features
    for feat, (desc, values) in binary_features.items():
        if feat in feature_dict:
            # Round to nearest integer since standardized values might not be exactly 0 or 1
            val = round(float(feature_dict[feat]))
            # Ensure val is either 0 or 1
            val = min(max(val, 0), 1)
            interpretations.append(f"{desc}: {values[val]}")
    
    # Process ordinal features - need to reverse standardization
    for feat, (desc, values) in ordinal_features.items():
        if feat in feature_dict:
            # For standardized values, we can only give an approximate interpretation
            std_val = feature_dict[feat]
            if std_val < -1.5:
                category = "Very Low"
            elif std_val < -0.5:
                category = "Low"
            elif std_val < 0.5:
                category = "Medium"
            elif std_val < 1.5:
                category = "High"
            else:
                category = "Very High"
            interpretations.append(f"{desc}: {category}")
    
    # Process continuous features
    for feat, desc in continuous_features.items():
        if feat in feature_dict:
            std_val = feature_dict[feat]
            if feat == "BMI":
                if std_val < -1.5:
                    category = "Underweight"
                elif std_val < -0.5:
                    category = "Normal"
                elif std_val < 0.5:
                    category = "Overweight"
                elif std_val < 1.5:
                    category = "Obese"
                else:
                    category = "Severely Obese"
                interpretations.append(f"{desc}: {category}")
            else:
                # For MentHlth and PhysHlth
                if std_val < -0.5:
                    days = "Few or No Days"
                elif std_val < 0.5:
                    days = "Some Days"
                elif std_val < 1.5:
                    days = "Many Days"
                else:
                    days = "Most Days"
                interpretations.append(f"{desc}: {days}")
    
    return interpretations

# 11. Show a few sample predictions
print("\nSample predictions (first 3 test samples):")
for i in range(3):
    idx = i
    features = X_test[:, idx]
    pred = y_pred_labels[idx]
    true = y_true_labels[idx]
    
    print(f"\nSample {i+1}:")
    print(f"  Predicted: {'Diabetes/Prediabetes' if pred == 1 else 'No Diabetes'}")
    print(f"  Actual: {'Diabetes/Prediabetes' if true == 1 else 'No Diabetes'}")
    print("  Features:")
    
    # Get human-readable interpretations
    interpretations = interpret_features(features, feature_names)
    for interp in interpretations:
        print(f"    {interp}")